#include <rtthread.h>

#include "lwip/sockets.h"

#define MAX_SERV 12
#define PORTFW_THREAD_NAME "portfw"
#if RT_THREAD_PRIORITY_MAX == 32
#define PORTFW_PRIORITY 17
#else
#define PORTFW_PRIORITY 170
#endif
#define PORTFW_THREAD_STACKSIZE 2048
#define BUFSIZE 128
#define max(x,y) ((x) > (y) ? (x) : (y))

struct portfwcb
{
    struct sockaddr_in srcaddr;
    struct sockaddr_in dstaddr;
};
struct clientcb
{
    struct clientcb *next;
    int socket;
    int fwsocket;
    struct sockaddr_in cliaddr;
    socklen_t clilen;
    char clibuf[BUFSIZE];
    char fwbuf[BUFSIZE];
    int clibuf_avail;
    int clibuf_written;
    int fwbuf_avail;
    int fwbuf_written;
};
static struct portfwcb sPortfwcb;
static struct clientcb *clientcb_list = 0;

static void close_client(struct clientcb *p_clientcb);
static int do_readcli(struct clientcb *p_clientcb);
static int do_readfw(struct clientcb *p_clientcb);
static int do_writecli(struct clientcb *p_clientcb);
static int do_writefw(struct clientcb *p_clientcb);

#define SHUT_CLI do {\
    if (p_clientcb->socket) {\
        rt_kprintf("Shutdown cli\n");\
        lwip_close(p_clientcb->socket); \
        p_clientcb->socket = -1; \
    }\
    } while(0)

#define SHUT_FW do {\
    if (p_clientcb->fwsocket) {\
        rt_kprintf("Shutdown fw\n");\
        lwip_close(p_clientcb->fwsocket); \
        p_clientcb->fwsocket = -1; \
    }\
    } while(0)

static void portfw_thread(void *arg)
{
    int listenfd;
    fd_set readset;
    fd_set writeset;
    fd_set errset;
    int i, maxfdp1 = 1;
    int clicnt = 0;
    //int keepAlive = 1;
    //int keepAliveSec = 2;
    struct clientcb *p_clientcb;
    //struct linger clilinger = {1, 2};
    listenfd = lwip_socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    LWIP_ASSERT("portfw_thread(): Socket create failed.", listenfd >= 0);
    if (lwip_bind(listenfd, (struct sockaddr*)&sPortfwcb.srcaddr, sizeof(struct sockaddr_in)) == -1)
        LWIP_ASSERT("portfw_thread(): Socket bind failed.", 0);

    if (lwip_listen(listenfd, MAX_SERV) == -1)
        LWIP_ASSERT("portfw_thread(): Listen failed.", 0);

    for (;;)
    {
        clicnt = 0;
        maxfdp1 = listenfd + 1;
        FD_ZERO(&readset);
        FD_ZERO(&writeset);
        FD_ZERO(&errset);
        FD_SET(listenfd, &readset);
        for (p_clientcb = clientcb_list; p_clientcb; p_clientcb = p_clientcb->next)
        {
            ++clicnt;
            if (p_clientcb->socket < 0 || (p_clientcb->fwsocket < 0 && p_clientcb->fwbuf_avail - p_clientcb->fwbuf_written == 0))
            {
                close_client(p_clientcb);
                break;
            }
            //rt_kprintf("fd: %d,%d\n", p_clientcb->socket, p_clientcb->fwsocket);
            maxfdp1 = max(maxfdp1, (p_clientcb->socket + 1));
            maxfdp1 = max(maxfdp1, (p_clientcb->fwsocket + 1));
            //rt_kprintf("maxfdp1: %d\n", maxfdp1);
            FD_SET(p_clientcb->socket, &readset);
            FD_SET(p_clientcb->fwsocket, &readset);
            FD_SET(p_clientcb->socket, &writeset);
            FD_SET(p_clientcb->fwsocket, &writeset);
            FD_SET(p_clientcb->socket, &errset);
            FD_SET(p_clientcb->fwsocket, &errset);
        }
        //rt_kprintf("Num cli: %d, maxfdp1: %d\n", clicnt, maxfdp1);

        i = lwip_select(maxfdp1, &readset, &writeset, &errset, 0);

        if (i == 0) continue;

        if (FD_ISSET(listenfd, &readset))
        {
            p_clientcb = (struct clientcb*)rt_calloc(1, sizeof(struct clientcb));
            if (p_clientcb)
            {
                p_clientcb->socket = lwip_accept(listenfd, (struct sockaddr*)&p_clientcb->cliaddr, &p_clientcb->clilen);
                if (p_clientcb->socket < 0)
                    rt_free(p_clientcb);
                else
                {
                    //lwip_setsockopt(p_clientcb->socket, SOL_SOCKET, SO_KEEPALIVE, (void*)&keepAlive, sizeof(keepAlive));
                    //lwip_setsockopt(p_clientcb->socket, IPPROTO_TCP, TCP_KEEPALIVE, (void*)&keepAliveSec, sizeof(keepAliveSec));
                    //lwip_setsockopt(p_clientcb->socket, SOL_SOCKET, SO_LINGER, &clilinger, sizeof(clilinger));
                    rt_kprintf("Creating fwsock\n");
                    p_clientcb->fwsocket = lwip_socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
                    if (p_clientcb->fwsocket < 0)
                    {
                        lwip_close(p_clientcb->socket);
                        rt_free(p_clientcb);
                    }
                    else
                    {
                        rt_kprintf("Connecting fwsock\n");
                        if (connect(p_clientcb->fwsocket, (struct sockaddr*)&(sPortfwcb.dstaddr), sizeof(sPortfwcb.dstaddr)) == -1)
                        {
                            rt_kprintf("Fail to connect fwsock\n");
                            lwip_close(p_clientcb->socket);
                            rt_free(p_clientcb);
                            continue;
                        }
                        p_clientcb->clibuf_avail = p_clientcb->clibuf_written = 0;
                        p_clientcb->fwbuf_avail = p_clientcb->fwbuf_written = 0;
                        p_clientcb->next = clientcb_list;
                        clientcb_list = p_clientcb;
                    }
                }
            }
            else
            {
                /* No memory to accept connection. Just accept and then close */
                int sock;
                struct sockaddr cliaddr;
                socklen_t clilen;
                rt_kprintf("No memory for new connection\n");
                sock = lwip_accept(listenfd, &cliaddr, &clilen);
                if (sock >= 0)
                    lwip_close(sock);
            }
        }

        /* Go through list of connected clients and process data */
        for (p_clientcb = clientcb_list; p_clientcb; p_clientcb = p_clientcb->next)
        {
            if (FD_ISSET(p_clientcb->socket, &errset))
            {
                rt_kprintf("client socket err\n");
                close_client(p_clientcb);
                break;
            }
            if (FD_ISSET(p_clientcb->fwsocket, &errset))
            {
                rt_kprintf("fw socket err\n");
                SHUT_FW;
                if (p_clientcb->fwbuf_avail - p_clientcb->fwbuf_written == 0)
                {
                    close_client(p_clientcb);
                    break;
                }
            }
            if (p_clientcb->socket > 0 && FD_ISSET(p_clientcb->socket, &readset))
            {
                rt_kprintf("cli readset\n");
                do_readcli(p_clientcb);
            }
            if (p_clientcb->fwsocket > 0 && FD_ISSET(p_clientcb->fwsocket, &readset))
            {
                rt_kprintf("fw readset\n");
                do_readfw(p_clientcb);
            }
            if (p_clientcb->socket > 0 && FD_ISSET(p_clientcb->socket, &writeset))
            {
                //rt_kprintf("cli writeset\n");
                if (p_clientcb->fwbuf_avail - p_clientcb->fwbuf_written > 0)
                    do_writecli(p_clientcb);
            }
            if (p_clientcb->fwsocket > 0 && FD_ISSET(p_clientcb->fwsocket, &writeset))
            {
                //rt_kprintf("fw writeset\n");
                if (p_clientcb->clibuf_avail - p_clientcb->clibuf_written > 0)
                    do_writefw(p_clientcb);
            }

            if (p_clientcb->clibuf_written == p_clientcb->clibuf_avail)
            {
                //rt_kprintf("reseting clibuf\n");
                p_clientcb->clibuf_written = p_clientcb->clibuf_avail = 0;
            }
            if (p_clientcb->fwbuf_written == p_clientcb->fwbuf_avail)
            {
                //rt_kprintf("reseting fwbuf\n");
                p_clientcb->fwbuf_written = p_clientcb->fwbuf_avail = 0;
            }

            if (p_clientcb->socket < 0 && p_clientcb->clibuf_avail - p_clientcb->clibuf_written == 0)
            {
                SHUT_FW;
            }
            if (p_clientcb->fwsocket < 0 && p_clientcb->fwbuf_avail - p_clientcb->fwbuf_written == 0)
            {
                SHUT_CLI;
            }

            if (p_clientcb->socket < 0 && p_clientcb->fwsocket < 0)
            {
                close_client(p_clientcb);
                break;
            }
        }
        //rt_kprintf("loop restart here %d\n", i);
    }
}

static void close_client(struct clientcb *p_clientcb)
{
    struct clientcb *p_search_clientcb;
    if (p_clientcb->socket > 0)
        lwip_close(p_clientcb->socket);
    if (p_clientcb->fwsocket > 0)
        lwip_close(p_clientcb->fwsocket);

    if (clientcb_list == p_clientcb)
        clientcb_list = p_clientcb->next;
    else
    {
        for (p_search_clientcb = clientcb_list; p_search_clientcb; p_search_clientcb = p_search_clientcb->next)
        {
            if (p_search_clientcb->next == p_clientcb)
            {
                p_search_clientcb->next = p_clientcb->next;
                break;
            }
        }
    }

    rt_free(p_clientcb);
}

static int do_readcli(struct clientcb *p_clientcb)
{
    int readcount;
    // If no more buf to store, just return
    if (BUFSIZE - p_clientcb->clibuf_avail == 0)
        return 0;

    readcount = lwip_read(p_clientcb->socket, p_clientcb->clibuf, BUFSIZE - p_clientcb->clibuf_avail);
    if (readcount < 1)
    {
        rt_kprintf("Cli: read failed %d\n", readcount);
        SHUT_CLI;
    }
    else
    {
        rt_kprintf("Cli: recv %d\n", readcount);
        p_clientcb->clibuf_avail += readcount;
        rt_kprintf("Cli: avail %d, written %d\n", p_clientcb->clibuf_avail, p_clientcb->clibuf_written);
    }

    rt_kprintf("Leaving do_readcli()\n");
    return 0;
}

static int do_readfw(struct clientcb *p_clientcb)
{
    int readcount;
    // If no more buf to store, just return
    if (BUFSIZE - p_clientcb->fwbuf_avail == 0)
        return 0;

    readcount = lwip_read(p_clientcb->fwsocket, p_clientcb->fwbuf, BUFSIZE - p_clientcb->fwbuf_avail);
    if (readcount < 1)
    {
        rt_kprintf("Fw: read failed %d\n", readcount);
        SHUT_FW;
    }
    else
    {
        rt_kprintf("Fw: recv %d\n", readcount);
        p_clientcb->fwbuf_avail += readcount;
        rt_kprintf("Fw: avail %d, written %d\n", p_clientcb->fwbuf_avail, p_clientcb->fwbuf_written);
    }

    rt_kprintf("Leaving do_readfw()\n");
    return 0;
}

static int do_writecli(struct clientcb *p_clientcb)
{
    int writecount;
    if (p_clientcb->fwbuf_avail - p_clientcb->fwbuf_written == 0)
        return 0;
    writecount = lwip_write(p_clientcb->socket, p_clientcb->fwbuf + p_clientcb->fwbuf_written, p_clientcb->fwbuf_avail - p_clientcb->fwbuf_written);
    if (writecount < 1)
    {
        rt_kprintf("Cli: write failed %d\n", writecount);
        SHUT_CLI;
    }
    else
    {
        rt_kprintf("Cli: send %d\n", writecount);
        p_clientcb->fwbuf_written += writecount;
        rt_kprintf("Cli send: avail %d, written %d\n", p_clientcb->fwbuf_avail, p_clientcb->fwbuf_written);
    }
    rt_kprintf("Leaving do_writecli()\n");
    return 0;
}

static int do_writefw(struct clientcb *p_clientcb)
{
    int writecount;
    if (p_clientcb->clibuf_avail - p_clientcb->clibuf_written == 0)
        return 0;

    writecount = lwip_write(p_clientcb->fwsocket, p_clientcb->clibuf + p_clientcb->clibuf_written, p_clientcb->clibuf_avail - p_clientcb->clibuf_written);
    if (writecount < 1)
    {
        rt_kprintf("Fw: write failed %d\n", writecount);
        SHUT_FW;
    }
    else
    {
        rt_kprintf("Fw: send %d\n", writecount);
        p_clientcb->clibuf_written += writecount;
        rt_kprintf("Fw send: avail %d, written %d\n", p_clientcb->clibuf_avail, p_clientcb->clibuf_written);
    }
    rt_kprintf("Leaving do_writefw()\n");
    return 0;
}

void portfw_init(int srcPort, char *adstAddr, int dstPort)
{
    struct ip_addr dstAddr;
    rt_thread_t portfw;
    struct _ip_addr
    {
        rt_uint8_t addr0, addr1, addr2, addr3;
    } *addr;
    memset(&sPortfwcb, 0, sizeof(sPortfwcb));
    if (inet_aton(adstAddr, (struct in_addr*)&dstAddr) == 0)
        return;
    sPortfwcb.dstaddr.sin_len = sizeof(struct sockaddr_in);
    sPortfwcb.dstaddr.sin_family = AF_INET;
    sPortfwcb.dstaddr.sin_addr.s_addr = dstAddr.addr;
    sPortfwcb.dstaddr.sin_port = htons(dstPort);

    sPortfwcb.srcaddr.sin_len = sizeof(struct sockaddr_in);
    sPortfwcb.srcaddr.sin_family = AF_INET;
    sPortfwcb.srcaddr.sin_addr.s_addr = htonl(INADDR_ANY);
    sPortfwcb.srcaddr.sin_port = htons(srcPort);
    portfw = rt_thread_create(PORTFW_THREAD_NAME,
        portfw_thread, RT_NULL,
        PORTFW_THREAD_STACKSIZE,
        PORTFW_PRIORITY, 5);
    if (portfw != RT_NULL)
        rt_thread_startup(portfw);
}

#ifdef RT_USING_FINSH
#include <finsh.h>
void portfw(int srcPort, char *adstAddr, int dstPort)
{
    portfw_init(srcPort, adstAddr, dstPort);
}
FINSH_FUNCTION_EXPORT(portfw, forward [srcPort] to [adstAddr] [dstPort]);
#endif
